#include "parser/MLIRToyVisitor.h"


MLIRToyVisitor::MLIRToyVisitor(std::string filename, mlir::MLIRContext& context)
    : builder(&context), filename(filename){
        this->theModule = mlir::ModuleOp::create(builder.getUnknownLoc());
}

// add symbol
void MLIRToyVisitor::addSymbol(llvm::StringRef name, mlir::Value theVal) {
    this->symbolTable.insert(std::make_pair(name, theVal));
}

void MLIRToyVisitor::printSymbol() {
    std::string s;
    for(auto p : this->symbolTable){
        s.append(p.first);
        s.append(", ");
    }
    std::cerr << s << std::endl;
}

mlir::ModuleOp MLIRToyVisitor::getModule(){
    return this->theModule;
}

mlir::Location MLIRToyVisitor::loc(int line, int row) {
    return mlir::FileLineColLoc::get(builder.getStringAttr(filename), line, row);
}

std::any MLIRToyVisitor::visitModule(toyParser::ModuleContext *ctx) {
    this->builder.setInsertionPointToEnd(theModule.getBody());
    return visitChildren(ctx);
}

std::any MLIRToyVisitor::visitStmt(toyParser::StmtContext *ctx) {
    return visitChildren(ctx);
}

std::any MLIRToyVisitor::visitVar_assign(toyParser::Var_assignContext *ctx) {
    auto varName = ctx->Identifier()->getText();
    auto value = std::any_cast<mlir::Value>(this->visit(ctx->expression()));
    if(this->symbolTable.find(varName) != this->symbolTable.end()){
        this->symbolTable[varName] = value;
    }else{
        this->addSymbol(varName, value);
    }
    
    return nullptr;
}

std::any MLIRToyVisitor::visitExpression(toyParser::ExpressionContext *ctx) {
    if(!ctx->prim_expr().empty()){
        mlir::Value theVal = std::any_cast<mlir::Value>(this->visit(ctx->prim_expr(0)));
        if(ctx->prim_expr().size() == 2){
            auto anotherVal = std::any_cast<mlir::Value>(this->visit(ctx->prim_expr(1)));
            auto location = this->loc(ctx->start->getLine(), ctx->start->getCharPositionInLine());
            theVal = this->builder.create<mlir::toy::AddOp>(location, theVal, anotherVal);
        }
        return theVal;
    }else{
        auto location = this->loc(ctx->Number()->getSymbol()->getLine(), ctx->Number()->getSymbol()->getCharPositionInLine());
        int num = std::stoi(ctx->Number()->getText());
        mlir::Value theVal = this->builder.create<mlir::toy::ConstantOp>(location, num);
        return theVal;
    }
}

std::any MLIRToyVisitor::visitPrint_op(toyParser::Print_opContext *ctx) {
    auto theVal = std::any_cast<mlir::Value>(this->visit(ctx->prim_expr()));
    auto location = this->loc(ctx->PRINT()->getSymbol()->getLine(), ctx->PRINT()->getSymbol()->getCharPositionInLine());
    this->builder.create<mlir::toy::PrintOp>(location, theVal);
    return nullptr;
}

std::any MLIRToyVisitor::visitPrim_expr(toyParser::Prim_exprContext *ctx) {
    if(ctx->Identifier()){
        auto varName = ctx->Identifier()->getText();
        if(this->symbolTable.count(varName) == 0){ // throw error
            std::cerr << "unknow identifider: " << varName << std::endl;
            throw "error";
            return nullptr;
        }else{
            return this->symbolTable[varName];
        }
    }else{
        auto location = this->loc(ctx->Number()->getSymbol()->getLine(), ctx->Number()->getSymbol()->getCharPositionInLine());
        int num = std::stoi(ctx->Number()->getText());
        mlir::Value theVal = this->builder.create<mlir::toy::ConstantOp>(location, num);
        return theVal;
    }
}